%% plot 3D band in qx, qy space
% here h is the diagonal term
% Bz/sqrt(2) = h
% h=[0,0.4,1/sqrt(2),2];
h=0.4;
r_sweep=linspace(0,sqrt(2),100);
alpha=0;
beta_sweep=linspace(0,2*pi,100);
phi=0;

for hlph=1:length(h)
    E=zeros(length(beta_sweep),length(r_sweep),3);
    for hlpr=1:length(r_sweep)
        r=r_sweep(hlpr);
        for hlpa=1:length(beta_sweep)
            beta=beta_sweep(hlpa);
            H0_tmp=Ham_detune(alpha,beta,0,0)*r+diag([h(hlph),0,-h(hlph)]);
            E_tmp=eig(H0_tmp);
            E_tmp=sort(E_tmp);
            E(hlpa,hlpr,:)=E_tmp;
        end
    end
    
    [R, PHI] = meshgrid(r_sweep,beta_sweep);
    figure
    surf(R.*cos(PHI),R.*sin(PHI),E(:,:,1),'EdgeColor','none');
    hold on
    surf(R.*cos(PHI),R.*sin(PHI),E(:,:,2),'EdgeColor','none');
    s=surf(R.*cos(PHI),R.*sin(PHI),E(:,:,3),'EdgeColor','none');
    view(-31.8,7.8)
    
    plot3(sqrt(2)*h(hlph)*cos(beta_sweep),sqrt(2)*h(hlph)*sin(beta_sweep),min(min(E(:,:,2)))*ones(size(beta_sweep)),'k--','LineWidth',2)
    
    colormap(jet)
    xlabel('q_x','FontSize',18)
    ylabel('q_y','FontSize',18)
    set(gca,'FontSize',18)
    xlim([-1,1])
    ylim([-1,1])
    set(gca,'ZTick',[])
    zlabel('Energy (a.u.)','FontSize',18)
end

% saveas(gcf,'FigS2A','epsc')

%% plot 3D band in qx, qz space

% h=[0,0.4,1/sqrt(2),2];
h=0.4;
r_sweep=linspace(0,sqrt(2),100);
alpha_sweep=linspace(0,2*pi,100);
beta=0;
phi=0;

for hlph=1:length(h)
    E=zeros(length(alpha_sweep),length(r_sweep),3);
    for hlpr=1:length(r_sweep)
        r=r_sweep(hlpr);
        for hlpa=1:length(alpha_sweep)
            alpha=alpha_sweep(hlpa);
            H0_tmp=Ham_detune(alpha,beta,0,0)*r+diag([h(hlph),0,-h(hlph)]);
            E_tmp=eig(H0_tmp);
            E_tmp=sort(E_tmp);
            E(hlpa,hlpr,:)=E_tmp;
        end
    end
    
    [R, ALPHA] = meshgrid(r_sweep,alpha_sweep);
    figure
    surf(R.*cos(ALPHA),R.*sin(ALPHA),E(:,:,1),'EdgeColor','none');
    hold on
    surf(R.*cos(ALPHA),R.*sin(ALPHA),E(:,:,2),'EdgeColor','none');
    s=surf(R.*cos(ALPHA),R.*sin(ALPHA),E(:,:,3),'EdgeColor','none');
    view(-44.1,5.98)
    
    colormap(jet)
    xlabel('q_x','FontSize',18)
    ylabel('q_z','FontSize',18)
    set(gca,'FontSize',18)
    xlim([-1,1])
    ylim([-1,1])
    set(gca,'ZTick',[])
    zlabel('Energy (a.u.)','FontSize',18)
end

% saveas(gcf,'FigS2B','epsc')


%% 1D plot with Bz, lambda4
% qy=qw=0, project to qx

qy=0;qw=0;
qx=linspace(-1,1,501);qz=qx;
num=length(qx);

E_out=zeros(num,num,3);
for hlpx=1:num
    for hlpz=1:num
        Ham_tmp=Ham_4D(qx(hlpx),qy,qz(hlpz),qw)+[0,0,1;0,0,0;1,0,0]*0.2*2+diag([1,0,-1])*0.2/sqrt(2)*2;
        E=sort(eig(Ham_tmp));
        
        E_out(hlpx,hlpz,1)=E(1);
        E_out(hlpx,hlpz,2)=E(2);
        E_out(hlpx,hlpz,3)=E(3);
    end
end

figure
cc=colormap(piratepal(9,1));
hold on
plot(qx,max(E_out(:,:,1),[],2),'.-','Color',cc(1,:))
plot(qx,min(E_out(:,:,2),[],2),'.','Color',cc(2,:))
plot(qx,max(E_out(:,:,2),[],2),'.','Color',cc(2,:))
plot(qx,min(E_out(:,:,3),[],2),'.-','Color',cc(3,:))
xlabel('q_x','FontSize',16)
ylabel('Energy (a.u.)','FontSize',18)
set(gca,'YTick',[],'FontSize',18)

% saveas(gcf,'FigS3A','epsc')


%% 1D plot with Bz, lambda5
% qy=qw=0, project to qx

qy=0;qw=0;
qx=linspace(-1,1,501);qz=qx;
num=length(qx);

E_out=zeros(num,num,3);
for hlpx=1:num
    for hlpz=1:num
        Ham_tmp=Ham_4D(qx(hlpx),qy,qz(hlpz),qw)+[0,0,-1i;0,0,0;1i,0,0]*0.2*2+diag([1,0,-1])*0.2/sqrt(2)*2;
        E=sort(eig(Ham_tmp));
        
        E_out(hlpx,hlpz,1)=E(1);
        E_out(hlpx,hlpz,2)=E(2);
        E_out(hlpx,hlpz,3)=E(3);
    end
end

figure
cc=colormap(piratepal(9,1));
hold on
plot(qx,max(E_out(:,:,1),[],2),'.-','Color',cc(1,:))
plot(qx,min(E_out(:,:,2),[],2),'.','Color',cc(2,:))
plot(qx,max(E_out(:,:,2),[],2),'.','Color',cc(2,:))
plot(qx,min(E_out(:,:,3),[],2),'.-','Color',cc(3,:))
xlabel('q_x','FontSize',16)
ylabel('Energy (a.u.)','FontSize',18)
set(gca,'YTick',[],'FontSize',18)

% saveas(gcf,'FigS3B','epsc')


function Ham_out = Ham_detune(alpha,theta,phi,rd) %
Ham_out=[0,cos(alpha)*exp(-1i*theta),0;cos(alpha)*exp(1i*theta),0,sin(alpha)*exp(1i*phi);...
    0,sin(alpha)*exp(-1i*phi),0];
Ham_out=Ham_out+diag([rd,0,-rd]);
end


function Ham_out = Ham_4D(qx,qy,qz,qw) %
Ham_out=[0,qx-1i*qy,0;qx+1i*qy,0,qz+1i*qw;...
    0,qz-1i*qw,0];
end

function E_out = E_spectrum(qx,qz,Bz,L4)
N1=1/3*(Bz.^2+L4.^2+qx.^2+qz.^2);
N32=1/2*Bz.*(qx.^2-qz.^2);
X=(N32+sqrt((N32.^2+L4.*qz.*qx).^2-N1.^3)+L4.*qx.*qz).^(1/3);

E_tmp1=X+N1./X;
E_tmp2=-sqrt(3)/2*1i*(X-N1./X)-1/2*(X+N1./X);
E_out=zeros(size(E_tmp1,1),size(E_tmp1,2),2);
E_out(:,:,1)=E_tmp1;
E_out(:,:,2)=E_tmp2;

end